#!/usr/bin/env python
# Module pcur test version 0.0.1
# Released to the public domain Aug,15,2006, 
# by wanyancan (wanyancan@gmail.com).
# License: GNU General Public License 2.0

# Provided as-is; use at your own risk; no warranty; no promises; enjoy!
# Use the Sources, Obiwan!
#                 --- Yoda

import pycurl, StringIO, traceback
import sys, re, os
import logging
#logging.basicConfig(format='%(asctime)s %(levelname)8s %(message)s', filename='/tmp/logfile.txt', filemode='w', level=logging.DEBUG)
logging.basicConfig(format='%(asctime)s %(levelname)8s %(message)s\n===debug-end===', level=logging.WARNING)

timeout = 30
numblock = 6

def pr_version():
    print """Pycurl version %s"""%pycurl.version_info()[1]
    
def errurl(err=None):
    print """Usage: pcur_test.py <url>
    """
    sys.exit(1)

def infrange(prefix, start=1):
    yield prefix
    while True:
        yield prefix+"."+str(start)
        start+=1

def tryoutfile(filename, overwrite = True, appendmod = False):
    rightpos = filename.rfind('?')
    if rightpos != -1:
        filename = filename[:rightpos]
    for ftry in infrange(filename):
        if not os.access(ftry, os.F_OK):
            return open(ftry, 'wb')
        elif appendmod:
            return (open(ftry, 'ab'),os.path.getsize(ftry))

def progress_callback(dltotal, dlnow, ultotal, ulnow):
    return False

def getsize( url):
    curl=pycurl.Curl()
    curl.setopt(pycurl.HEADER, True)
    curl.setopt(pycurl.NOBODY, True)
    curl.setopt(pycurl.URL, url)
    curl.setopt(pycurl.TIMEOUT, timeout)

    b = StringIO.StringIO()
    curl.setopt(pycurl.WRITEFUNCTION, b.write)
    curl.perform()
    logging.debug(b.getvalue())
    try:
        size = int(re.findall("Content-Length: (\d+)", b.getvalue())[0])
    except:
        size = -1
    return size
    
curlpool = []
filepool = []

# file[pos] that start <= pos < end should be transfered
def gen_curl(url,start=-1, end=-1, filesize=-1):
    curl = pycurl.Curl()
    curl.setopt(pycurl.URL, url)
    curl.url = url
    if start!=-1:
        # the RANGE option will get file[pos] what l<=pos<=r
        if end>filesize and filesize!=-1:
            end = filesize
        curl.setopt(pycurl.RANGE, str(start)+"-"+str(end-1))
        curl.range=(start, end)
        logging.debug("getting %d-%d(%d) of file: %s", start, end-1,filesize, url)
        curl.filename=getfilename(url)
        f = tryoutfile("%s$%d$%d"%(getfilename(url),start,end), appendmod=True)
        if isinstance(f, tuple):
            start += f[1]
            curl.setopt(pycurl.RANGE, str(start)+"-"+str(end-1))
            curl.range=(start, end)
            logging.debug("resuming %d-%d(%d) of file: %s", start, end-1,filesize, url)
            f = f[0]
        curl.fp = f
        filepool.append(f)
        curl.setopt(pycurl.WRITEDATA, f)

    else:
        curl.setopt(pycurl.NOPROGRESS, False)
        curl.setopt(pycurl.PROGRESSFUNCTION, progress_callback)
    return curl

def getfilename(url):
    if url.endswith("/"):
        return "index.html"
    else:
        try:
            return url[url.rindex('/')+1:]
        except ValueError, err:
            errurl(err)

def _showcurl_status(curlhandles):
    _s = "\x1b\x5b"
    sys.stdout.write(_s+"%dA"%(len(curlhandles)+1))
    totspeed = 0.0
    totget = 0
    totneed = 0
    for i in range(len(curlhandles)):
        curl = curlhandles[i]
        downloaded = curl.getinfo(pycurl.SIZE_DOWNLOAD)
        totget+=downloaded
        speed = curl.getinfo(pycurl.SPEED_DOWNLOAD)
        totspeed+=speed
        start, end = curl.range[:]
        totneed += end-start
        if end==start:
            print "this part has already finished"
        else:
            print "%d of part(%d-%d), %d%% finished (%.2fkB/s)"%(int(downloaded), start, end-1, int(downloaded/(end-start)*100), speed/1000.0)
    if totneed==0:
        print "all parts have been finished"
    else:
        print "total %d of %d(%d%%), at rate %.3fkB/s"%(totget, totneed,int(100.0*totget/totneed), totspeed/1000.0)

def end_main(bNormal):
    fname = []
    for f in filepool:
        f.flush()
        fname.append(f.name)
        f.close()
    if bNormal:
        ftmpout = open(fname[0], 'ab')
        for i in range(1, len(fname)):
            ftmpin = open(fname[i], 'rb')
            for line in ftmpin.xreadlines():
                ftmpout.write(line)
            ftmpin.close()
            os.remove(fname[i])
        ftmpout.close()
        newname = fname[0][:fname[0].find('$')]
        os.rename(fname[0], newname)
        print "written to file: %s"%newname
    else:
        print "user aborted."
    pycurl.global_cleanup()

if __name__=='__main__':
    import doctest, sys
    doctest.testmod(sys.modules[__name__])

    pr_version()
    if len(sys.argv) < 2:
        errurl()

    url = sys.argv[1]

    filename = getfilename(url)
    logging.debug('outfile written to %s', filename)

    pycurl.global_init(pycurl.GLOBAL_ALL) # GLOBAL_ALL must be set in normal

    filesize = getsize(url)
    if filesize ==-1:  # length not known
        c = gen_curl(url)
        outfile = tryoutfile(filename)        
        c.setopt(pycurl.WRITEFUNCTION, outfile.write)
        c.perform()
    else:
        blocksize = filesize/numblock +1
        for p_sta,p_end in [(x,x+blocksize) for x in xrange(0, filesize, blocksize)]:
            curlpool.append(gen_curl(url, p_sta, p_end, filesize))
        m = pycurl.CurlMulti()
        m.handles = []
        for c in curlpool:
            m.add_handle(c)
            m.handles.append(c)
        print "starting all curl in pool..."
        print "getting url: %s"%url
        try:
            while True:
                ret, num_handles = m.perform()
                if ret != pycurl.E_CALL_MULTI_PERFORM: break
            print '\n'*len(m.handles)
            _showcurl_status(m.handles)            
            while num_handles:
                ret = m.select(1.0)
                if ret == -1:  continue

                while 1:
                    ret, num_handles = m.perform()
                    if ret != pycurl.E_CALL_MULTI_PERFORM: break
                _showcurl_status(m.handles)
            end_main(bNormal=True)
        except KeyboardInterrupt:
            end_main(bNormal=False)
        except SystemExit:
            end_main(bNormal=False)
